{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nOCkwFB0WHCV"
   },
   "source": [
    "##### Copyright 2021 The OpenFermion Developers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "G24eMxuEWA04"
   },
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "795sU01of3pF"
   },
   "source": [
    "# Simulating the Fermi-Hubbard model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ykV-VNQ_VI4h"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://quantumai.google/openfermion/fqe/tutorials/fermi_hubbard\"><img src=\"https://quantumai.google/site-assets/images/buttons/quantumai_logo_1x.png\" />View on QuantumAI</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/quantumlib/OpenFermion/blob/master/docs/fqe/tutorials/fermi_hubbard.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/colab_logo_1x.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/quantumlib/OpenFermion/blob/master/docs/fqe/tutorials/fermi_hubbard.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/github_logo_1x.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/OpenFermion/docs/fqe/tutorials/fermi_hubbard.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/download_icon_1x.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_Yg1R294f9QU"
   },
   "source": [
    "This notebook shows how to simulate the one-dimensional Fermi-Hubbard Hamiltonian\n",
    "\n",
    "$$\n",
    "H = - J \\sum_{j = 1}^{L - 1} \\sum_{\\sigma \\in \\{ \\uparrow, \\downarrow \\}} c_{j, \\sigma}^\\dagger c_{j + 1, \\sigma} + \\text{h.c.} + U \\sum_{j} n_{j\\uparrow} n_{j\\downarrow}\n",
    "$$\n",
    "\n",
    "using FQE. Here $j = 1, ..., L$ denotes site/orbital and $\\sigma \\in \\{ \\uparrow, \\downarrow \\}$ denotes spin. By the end of the tutorial, we reproduce plots from the [Fermi-Hubbard experiment paper](https://arxiv.org/abs/2010.07965) and the corresponding [ReCirq tutorial](https://quantumai.google/cirq/experiments/fermi_hubbard/experiment_example)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PDSKaWoFsRTv"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import fqe\n",
    "except ImportError:\n",
    "    # TODO: Change to `pip install fqe --quiet` when FQE>0.1.0 is released.\n",
    "    !pip install git+https://github.com/quantumlib/OpenFermion-FQE --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rEssoRYDsTBK"
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "from itertools import product\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy import sparse\n",
    "from scipy.linalg import expm\n",
    "\n",
    "import openfermion as of\n",
    "import fqe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-TdjpB6Zh0y_"
   },
   "source": [
    "## Defining the Hamiltonian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lb4FCZa-e0fm"
   },
   "source": [
    "We first define the Fermi-Hubbard Hamiltonian using OpenFermion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "poyfiq0GQF1o"
   },
   "outputs": [],
   "source": [
    "\"\"\"Define the Hamiltonian.\"\"\"\n",
    "# Parameters.\n",
    "nsites = 4\n",
    "U = 2.0\n",
    "J = -1.0\n",
    "\n",
    "hubbard = of.fermi_hubbard(1, nsites, tunneling=-J, coulomb=U, periodic=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w-9K9H-xh4IT"
   },
   "source": [
    "## Time evolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dvswbGDUKT6q"
   },
   "source": [
    "In this section we show how to do time evolution with FQE. We start from a random initial state in the two-particle $S_z = 0$ sector to ignore the details of state preparation in the [Fermi-Hubbard experiment paper](https://arxiv.org/abs/2010.07965) for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KWbyc7ZXiSpw"
   },
   "outputs": [],
   "source": [
    "nele, sz = 2, 0\n",
    "init_wfn = fqe.Wavefunction([[nele, sz, nsites]])\n",
    "init_wfn.set_wfn(strategy=\"random\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TDOWVkSgAh5H"
   },
   "source": [
    "To do the time evolution, we can simply call the `fqe.Wavefunction.time_evolve` method with the Hubbard Hamiltonian and evolution time. When passing `time_evolve` a general fermion operator, FQE converts this into a dense Hamiltonian for time evolution. The evolution is accomplished by a Taylor expansion method.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pmj3lmSItPLu"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "e_time = 1.0\n",
    "true_evolved_fqe = init_wfn.time_evolve(e_time, hubbard)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LZLGYR82ewFs"
   },
   "source": [
    "We can compare the speed of this method to direct matrix exponentiation as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YjAKYrZq_gfQ"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert FQE wavefunction to an 2^n vector |𝛹⟩.\n",
    "init_cirq_wfn = fqe.to_cirq(init_wfn)\n",
    "\n",
    "# Exponentiate the (sparse) Hamiltonian U = exp(-i H t).\n",
    "unitary = expm(-1j * e_time * of.get_sparse_operator(hubbard))\n",
    "\n",
    "# Do the time evolution |𝛹'⟩ = U |𝛹⟩.\n",
    "true_evolved_cirq = unitary @ init_cirq_wfn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p_d7lu54gZF5"
   },
   "source": [
    "The next cell verifies both methods produce the same final wavefunction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "93bMs2yov-49"
   },
   "outputs": [],
   "source": [
    "fidelity = abs(fqe.vdot(true_evolved_fqe, fqe.from_cirq(true_evolved_cirq, thresh=1e-12)))**2\n",
    "assert np.isclose(fidelity, 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aXsPEYrlhV89"
   },
   "source": [
    "Even for a small number of sites, the FQE simulation is notably faster than the direct matrix exponentian simulation. This is because FQE evolves in only fixed particle and spin subspace and because the time evolution is simulated by series expansion of the matrix exponent. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Fwrhs1Whigaf"
   },
   "source": [
    "### Trotterization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9EuhssdinFxR"
   },
   "source": [
    "We can apply the Trotter formula and evolve the hopping Hamiltonian (one-body terms)\n",
    "\n",
    "$$\n",
    "H_{\\text{hop}} := - J \\sum_{j = 1}^{L - 1} \\sum_{\\sigma \\in \\{ \\uparrow, \\downarrow \\}} c_{j, \\sigma}^\\dagger c_{j + 1, \\sigma} + \\text{h.c.} \n",
    "$$\n",
    "\n",
    "and the charge-charge interaction (two-body terms)\n",
    "\n",
    "$$\n",
    "H_{\\text{cc}} := U \\sum_{j} n_{j\\uparrow} n_{j\\downarrow}\n",
    "$$\n",
    "\n",
    "separately for short time steps. We first split up the Hamiltonian into one-body and two-body terms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NDkrbIXph07d"
   },
   "outputs": [],
   "source": [
    "# Each site has two spins.\n",
    "nqubits = 2 * nsites\n",
    "\n",
    "# One-body (hopping) terms.\n",
    "one_body_terms = [\n",
    "    op + of.hermitian_conjugated(op) for op in (\n",
    "        of.FermionOperator(((i, 1), (i + 2, 0)), coefficient=J) for i in range(nqubits - 2)\n",
    "    )\n",
    "]\n",
    "\n",
    "# Two-body (charge-charge) terms.\n",
    "two_body_terms = [\n",
    "    of.FermionOperator(((i, 1), (i, 0), (i + 1, 1), (i + 1, 0)), coefficient=U)\n",
    "    for i in range(0, nqubits, 2)\n",
    "]\n",
    "\n",
    "# Verify this produces the same Hamiltonian from OpenFermion.\n",
    "assert sum(one_body_terms) + sum(two_body_terms) == hubbard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l_XqVYdY3cXI"
   },
   "source": [
    "The Hopping Hamiltonian is a $2^n \\times 2^n$ matrix, so one method for time evolution is directly exponentiating this matrix to apply $e^{- i t H_{\\text{hop}}}$ to the initial state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b9sA8quStEfA"
   },
   "outputs": [],
   "source": [
    "# Get the 2^n x 2^n matrix (as a sparse operator).\n",
    "sparse_hopping_matrix = of.get_sparse_operator(sum(one_body_terms))\n",
    "\n",
    "# Exponentiate the matrix.\n",
    "unitary = expm(-1j * e_time * sparse_hopping_matrix)\n",
    "\n",
    "# Time-evolve the initial state.\n",
    "evolved_cirq_wfn = unitary @ init_cirq_wfn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wNUBa8Tv5jIR"
   },
   "source": [
    "Now let's see how to do this in FQE. We first recognize the hopping Hamiltonian can be written as the following $n \\times n$ matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GL_tobpNZ4GP"
   },
   "outputs": [],
   "source": [
    "hopping_matrix = J * (np.diag([1] * (nsites - 1), k=1) + np.diag([1] * (nsites - 1), k=-1))\n",
    "print(hopping_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u9NruSyh66Zg"
   },
   "source": [
    "Using the unitary generated by this matrix, we can evolve the initial state using the function `evolve_fqe_givens` as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A_1oGB4e6HoL"
   },
   "outputs": [],
   "source": [
    "from fqe.algorithm.low_rank import evolve_fqe_givens\n",
    "\n",
    "umat = expm(-1j * e_time * hopping_matrix)\n",
    "evolved_wfn = evolve_fqe_givens(init_wfn, umat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q8dtdB7nMStR"
   },
   "source": [
    "We can verify the wavefunction evolved by FQE and the wavefunction evolved by direct matrix exponentiation are identical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "w_8xDNr9sdAZ"
   },
   "outputs": [],
   "source": [
    "fidelity = abs(fqe.vdot(evolved_wfn, fqe.from_cirq(evolved_cirq_wfn, thresh=1e-6)))**2\n",
    "assert np.isclose(fidelity, 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VeXD_qUcMfsA"
   },
   "source": [
    "Like the Hopping Hamiltonian, the charge-charge interaction is a $2^n \\times 2^n$ matrix, so we can again apply $e^{- i t H_{\\text{cc}}}$ to the initial state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jusAhg5Lbv-r"
   },
   "outputs": [],
   "source": [
    "# Get the 2^n x 2^n matrix (as a sparse operator).\n",
    "sparse_charge_charge_matrix = of.get_sparse_operator(sum(two_body_terms))\n",
    "\n",
    "# Exponentiate the matrix.\n",
    "unitary = expm(-1j * e_time * sparse_charge_charge_matrix)\n",
    "\n",
    "# Time-evolve the initial state.\n",
    "evolved_cirq_wfn = unitary @ init_cirq_wfn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XRjdKI_29ENl"
   },
   "source": [
    "As with the hopping Hamiltonian, the charge-charge term can be written as an $n \\times n$ matrix as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "extNYasKbm5k"
   },
   "outputs": [],
   "source": [
    "charge_charge_matrix = np.diag([U] * nsites)\n",
    "charge_charge_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sI9eV5NTcFWF"
   },
   "source": [
    "We use this to perform time-evolution with the `evolve_fqe_charge_charge_alpha_beta` function. This function applies\n",
    "\n",
    "$$\n",
    "\\exp\\left(-i t \\sum_{i,j} v_{i, j} n_{i, \\alpha} n_{j, \\beta} \\right)\n",
    "$$\n",
    "\n",
    "for any coefficients $v_{i, j}$ to the wavefunction. (One can easiliy check the charge-charge interaction $H_{\\text{cc}}$ can be written in this way.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RYjs8McNj-hM"
   },
   "outputs": [],
   "source": [
    "from fqe.algorithm.low_rank import evolve_fqe_charge_charge_alpha_beta\n",
    "\n",
    "evolved_fqe_wfn = evolve_fqe_charge_charge_alpha_beta(\n",
    "    init_wfn, \n",
    "    charge_charge_matrix, \n",
    "    e_time\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3Tc6YDeoM5l9"
   },
   "source": [
    "As before we can verify this produces the same evolved wavefunction as direct matrix exponentiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rIu45SonLFUT"
   },
   "outputs": [],
   "source": [
    "fidelity = abs(fqe.vdot(evolved_fqe_wfn, fqe.from_cirq(evolved_cirq_wfn, thresh=1e-6)))**2\n",
    "assert np.isclose(fidelity, 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IrbBq5MvvYXK"
   },
   "source": [
    "At this point we can simulate a specified number of Trotter steps as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oDqQWKzElOzZ"
   },
   "outputs": [],
   "source": [
    "# Set the number of Trotter steps.\n",
    "trotter_steps = 10\n",
    "\n",
    "# Evolution under the hopping Hamiltonian for a single Trotter step.\n",
    "umat = expm(-1j * hopping_matrix * (e_time / trotter_steps))\n",
    "\n",
    "# Simulate each Trotter step.\n",
    "current_wfn = copy.deepcopy(init_wfn)\n",
    "for _ in range(trotter_steps):\n",
    "    # Evolve the Hopping Hamiltonian.\n",
    "    current_wfn = evolve_fqe_givens(current_wfn,  u=umat)\n",
    "\n",
    "    # Evolve the charge-charge interaction.\n",
    "    current_wfn = evolve_fqe_charge_charge_alpha_beta(\n",
    "        current_wfn, \n",
    "        charge_charge_matrix, \n",
    "        e_time / trotter_steps,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3ARWCKnWpzeT"
   },
   "source": [
    "Since we already evolved the wavefunction without Trotterization, we can compare the Trotterized evolution to this below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8B8PAKTDpVox"
   },
   "outputs": [],
   "source": [
    "trotterized_fidelity = abs(fqe.vdot(true_evolved_fqe, current_wfn))**2\n",
    "print(\"Trotterized fidelity:\", trotterized_fidelity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XhRwkJJfqDz2"
   },
   "source": [
    "One can increase the number of Trotter steps to increase this fidelity. When reproducing experimental results below, we will use Trotterization as in the paper."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GMEIG8XSfELb"
   },
   "source": [
    "## Reproducing experimental results\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0ZG88i9yNjzP"
   },
   "source": [
    "We now reproduce results from the [Fermi-Hubbard experiment paper](https://arxiv.org/abs/2010.07965) using the time-evolution methods in FQE we have discussed above. The only remaining ingredient is the initial state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LrmWZy4JhNaQ"
   },
   "source": [
    "### Preparing the initial state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GIp-KUxXhPV_"
   },
   "source": [
    "The initial state in the Fermi-Hubbard experiment is the ground state of a non-interacting Hamiltonian with different onsite potentials. The spin-up sector and spin-down sector have no coupling. We create the two non-interacting components of this Hamiltonian below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TkoqgM8cgdam"
   },
   "outputs": [],
   "source": [
    "# Parameters used in the Fermi-Hubbard experiment.\n",
    "nsites = 8\n",
    "l_up = 4\n",
    "m_up = 4.5\n",
    "sigma_up = 1\n",
    "\n",
    "# For convenience.\n",
    "site_index = np.arange(1, nsites + 1)\n",
    "\n",
    "# Creating the spin-up and spin-down Hamiltonians.\n",
    "spin_up_ham = np.diag([-1.] * (nsites - 1), k=1) + np.diag([-1.] * (nsites - 1), k=-1)\n",
    "spin_up_ham += np.diag(-l_up * np.exp(-0.5 * (site_index - m_up)**2) / sigma_up**2)\n",
    "spin_down_ham = np.diag([-1.] * (nsites - 1), k=1) + np.diag([-1.] * (nsites - 1), k=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vWE6nsqv_xpX"
   },
   "source": [
    "We know our one-body Hamiltonian is diagonalized by a quadratic Hamiltonian that can be implemented via a Givens rotation network. We now occupy the two lowest states $b^\\dagger_{1a} b^\\dagger_{2a} b^\\dagger_{1b} b^\\dagger_{2b}|\\text{vac}\\rangle$ and then reverse-transform back to the original basis. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Mh-3UhwcqmEy"
   },
   "outputs": [],
   "source": [
    "# Initialize wavefunction to occupy the two lowest states.\n",
    "init_wfn = fqe.Wavefunction([[nsites // 2, 0, nsites]])\n",
    "init_wfn.set_wfn(strategy='hartree-fock')\n",
    "init_wfn.print_wfn()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2McrJDncyrHG"
   },
   "source": [
    "To transform back, we first diagonalize the spin-up and spin-down hamiltonians to get the eigenbasis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k0gffkxxM_AK"
   },
   "outputs": [],
   "source": [
    "_, v_up = np.linalg.eigh(spin_up_ham)\n",
    "_, v_down = np.linalg.eigh(spin_down_ham)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dS-9TA8QYQQt"
   },
   "source": [
    "And now rotate back to the original basis through a one-body transform implemented with the Givens rotation network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NZ87egHwq-kc"
   },
   "outputs": [],
   "source": [
    "from fqe.algorithm.low_rank import evolve_fqe_givens_sector\n",
    "\n",
    "\n",
    "init_wfn = evolve_fqe_givens_sector(init_wfn, v_up, sector='alpha')\n",
    "init_wfn = evolve_fqe_givens_sector(init_wfn, v_down, sector='beta')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gQBWCf7Qp3OA"
   },
   "source": [
    "Now that we have the initial state, we can compute the charge ($+$) and spin ($-$) density, respectively\n",
    "\n",
    "$$\n",
    "\\rho_{j}^{\\pm} = \\langle n_{j, \\uparrow} \\rangle \\pm \\langle n_{j, \\downarrow} \\rangle,\n",
    "$$\n",
    "\n",
    "at time $t = 0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "hnnUFLfC1w6E"
   },
   "outputs": [],
   "source": [
    "#@title Plotting.\n",
    "plt.rcParams['font.family'] = 'sans-serif'\n",
    "plt.rcParams['font.size'] = 14\n",
    "colors = ['#4285F4', '#EA4335']\n",
    "\n",
    "# Plot the data.\n",
    "opdm_a, opdm_b = init_wfn.sector((nsites // 2, 0)).get_spin_opdm()\n",
    "charge_density = np.diagonal(opdm_a).real + np.diagonal(opdm_b).real\n",
    "spin_density = np.diagonal(opdm_a).real - np.diagonal(opdm_b).real\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(6, 4))\n",
    "plt.plot(site_index, charge_density, \"-o\", color=colors[0], label=\"Charge\")\n",
    "plt.plot(site_index, spin_density, \"-o\", color=colors[1], label=\"Spin\")\n",
    "\n",
    "# Axes.\n",
    "ax.set_ylabel(r\"Charge density $\\rho^+$\", color=colors[0])\n",
    "ax_twin = ax.twinx()\n",
    "ax_twin.set_ylabel(r\"Spin density $\\rho^-$\", color=colors[1])\n",
    "\n",
    "ax.set_xticks(site_index)\n",
    "ax.set_xlabel(\"Site index\")\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w11Ajd13_lvd"
   },
   "source": [
    "One can check this plot reproduces Fig. 2(a) from the [Fermi-Hubbard experiment paper](https://arxiv.org/abs/2010.07965). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NxRG6_z19m_K"
   },
   "source": [
    "#### Compare to direct diagonalization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mg1H9oNgYedN"
   },
   "source": [
    "Here we quickly compare our previous procedure for creating the initial state to exact diagonalization. We will see they are identical."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "02ZeA9Iu1gRu"
   },
   "source": [
    "The cell below creates the initial non-interacting Hamiltonian to diagonalize in two steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Gp-VtmEEblxj"
   },
   "outputs": [],
   "source": [
    "# (1) Create Fermion operator.\n",
    "h0 = np.zeros((2 * nsites, 2 * nsites))\n",
    "h0[::2, ::2] = spin_up_ham\n",
    "h0[1::2, 1::2] = spin_down_ham\n",
    "\n",
    "fop = of.FermionOperator()\n",
    "for p, q in product(range(2 * nsites), repeat=2):\n",
    "    fop += of.FermionOperator(((p, 1), (q, 0)), coefficient=h0[p, q])\n",
    "\n",
    "fop_mat = of.get_sparse_operator(fop)\n",
    "\n",
    "# (2) Project Fermion operator.\n",
    "dim = 2**(2 * nsites)\n",
    "diag_val = []\n",
    "diag_pos = []\n",
    "\n",
    "for ii in range(dim):\n",
    "    ket = np.binary_repr(ii, width=2*nsites)\n",
    "    ket_a = list(map(int, ket[::2]))\n",
    "    ket_b = list(map(int, ket[1::2]))\n",
    "\n",
    "    if np.isclose(sum(ket_a) + sum(ket_b), 4) and np.isclose(sum(ket_a) - sum(ket_b), 0):\n",
    "        assert np.isclose(sum(ket_a), 2)\n",
    "        assert np.isclose(sum(ket_b), 2)\n",
    "        diag_val.append(1)\n",
    "        diag_pos.append(ii)\n",
    "\n",
    "\n",
    "proj_n = sparse.coo_matrix((diag_val, (diag_pos, diag_pos)), shape=(dim, dim))\n",
    "fop_mat = proj_n @ fop_mat @ proj_n  # Hamiltonian to diagonalize."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BqOEjlVfjCbF"
   },
   "source": [
    "Now that we have the Hamiltonian, we can diagonalize it to get the ground state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h-03wat3hca1"
   },
   "outputs": [],
   "source": [
    "evals, evecs = sparse.linalg.eigsh(fop_mat.real, k=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ku-5DZ1UPPLp"
   },
   "source": [
    "At this point we have the initial wavefunction as a $2^n$ element vector. Below we convert this to the FQE representation and check it is identical to the previous initial wavefunction we prepared."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_W7FrgmvjAuS"
   },
   "outputs": [],
   "source": [
    "init_wfn_from_diag = fqe.from_cirq(evecs[:, 0].flatten(), thresh=1.0E-12)\n",
    "\n",
    "# Check.\n",
    "fidelity = abs(fqe.vdot(init_wfn, init_wfn_from_diag)) ** 2\n",
    "assert np.isclose(fidelity, 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KwJFUU36kDDU"
   },
   "source": [
    "### Computing the time-evolved charge and spin density"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XF94rfwW3ZRN"
   },
   "source": [
    "We now time-evolve the initial state under the Fermi-Hubbard model. As in the paper, we use time step $\\tau = 0.3 \\hbar / J$ for a total simulation time of $16.5 \\hbar / J$. (Note that we use atomic units where $\\hbar = 1$.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3tMot2Llu7Fw"
   },
   "outputs": [],
   "source": [
    "# Set the time step and number of Trotter steps (using hbar = 1).\n",
    "dt = 0.3 \n",
    "trotter_steps = 55\n",
    "\n",
    "hopping_u = expm(-1j * dt * spin_down_ham)\n",
    "charge_u = np.diag([U] * nsites)\n",
    "\n",
    "opdms_alpha = [opdm_a]\n",
    "opdms_beta = [opdm_b]\n",
    "real_times = [0.0]\n",
    "\n",
    "current_time = 0.0\n",
    "final_wfn = init_wfn\n",
    "for tt in range(trotter_steps):\n",
    "    final_wfn = evolve_fqe_givens(final_wfn, hopping_u)\n",
    "    final_wfn = evolve_fqe_charge_charge_alpha_beta(final_wfn, charge_u, dt)\n",
    "    current_time += dt\n",
    "    opdm_a_current, opdm_b_current = final_wfn.sector((nsites // 2, 0)).get_spin_opdm()\n",
    "\n",
    "    # Store the opdms at each time step.\n",
    "    opdms_alpha.append(opdm_a_current)\n",
    "    opdms_beta.append(opdm_b_current)\n",
    "    real_times.append(current_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RK7RLmWiUhpT"
   },
   "source": [
    "Using the stored data, we compute the charge and spin density at various times below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q6X3ai09blf3"
   },
   "outputs": [],
   "source": [
    "# Times to plot the charge & spin density at.\n",
    "times = [0.0, 1.2, 1.8, 3.0]\n",
    "time_indices = [int(t / dt) for t in times]\n",
    "\n",
    "charge_densities = []\n",
    "spin_densities = []\n",
    "for idx in time_indices:\n",
    "    charge_densities.append(np.diagonal(opdms_alpha[idx] + opdms_beta[idx]).real)\n",
    "    spin_densities.append(np.diagonal(opdms_alpha[idx] - opdms_beta[idx]).real)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e4cv9BicdYaZ"
   },
   "source": [
    " We now plot these values to reproduce Fig. 1(a) from the Fermi-Hubbard paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "GYp7rNoabwVN"
   },
   "outputs": [],
   "source": [
    "#@title Plotting.\n",
    "fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(10, 5), sharex=True)\n",
    "\n",
    "for (i, ax) in enumerate(axes.flatten()):\n",
    "    ax.plot(site_index, charge_densities[i], \"-o\", color=colors[0])\n",
    "    ax_twin = ax.twinx()\n",
    "    ax_twin.plot(site_index, spin_densities[i], \"-o\", color=colors[1])\n",
    "\n",
    "    ax.set_ylim([0, 1.2])\n",
    "    ax_twin.set_ylim([-0.3, 0.6])\n",
    "    ax.tick_params(which=\"both\", direction=\"in\")\n",
    "    ax_twin.tick_params(which=\"both\", direction=\"in\")\n",
    "\n",
    "    if i >= 2:\n",
    "        ax.set_xlabel(\"Site index\")\n",
    "    if i % 2 == 0:\n",
    "        ax.set_ylabel(r\"Charge density $\\rho^+$\", color=colors[0])\n",
    "    if i % 2 == 1:\n",
    "        ax_twin.set_ylabel(r\"Spin density $\\rho^-$\", color=colors[1])\n",
    "    ax.text(1, 1, rf\"t = ${times[i]} \\hbar / J$\")\n",
    "\n",
    "plt.xticks(site_index)\n",
    "plt.tight_layout()\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XsVthwofkaLT"
   },
   "source": [
    "### Computing the charge and spin spread"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z0R1BCz6shA_"
   },
   "source": [
    "Last, we use this data to compute the charge and spin spread."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0wtvSe9kzQHh"
   },
   "outputs": [],
   "source": [
    "charge_spread = []\n",
    "spin_spread = []\n",
    "spread_func = np.abs(np.arange(1, nsites + 1) - (nsites + 1) / 2)\n",
    "\n",
    "for ii in range(len(opdms_alpha)):\n",
    "    charge_spread.append(np.sum(np.multiply(spread_func, np.diagonal(opdms_alpha[ii] + opdms_beta[ii]).real)))\n",
    "    spin_spread.append(np.sum(np.multiply(spread_func, np.diagonal(opdms_alpha[ii] - opdms_beta[ii]).real)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "2UHdd3Rt33LB"
   },
   "outputs": [],
   "source": [
    "#@title Plotting.\n",
    "fig, ax = plt.subplots(nrows=1, ncols=1)\n",
    "\n",
    "stop = 11\n",
    "ax.plot(real_times[:stop], charge_spread[:stop], \"-o\", color=colors[0])\n",
    "axtwin = ax.twinx()\n",
    "axtwin.plot(real_times[:stop], spin_spread[:stop], \"-o\", color=colors[1])\n",
    "\n",
    "ax.set_ylim([5, 10])\n",
    "axtwin.set_ylim([-2.5, 2.6]);\n",
    "\n",
    "ax.set_xticks(real_times[:stop])\n",
    "\n",
    "ax.set_xlabel(r\"Time $(\\hbar / J)$\")\n",
    "ax.set_ylabel(r\"Charge spread $\\kappa^+$\", color=colors[0])\n",
    "axtwin.set_ylabel(r\"Spin spread $\\kappa^-$\", color=colors[1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "arVpvshwU5Py"
   },
   "source": [
    "One can check this reproduces the $u = 0$ subplot from Fig. 2(b) in the Fermi-Hubbard paper."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "fermi_hubbard.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
